package tbsharding

import (
	"context"
	"fmt"
	"gitee.com/youbeiwuhuan/gin-practice/go-gorm2-sharding/entity"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
	"gorm.io/sharding"
	"strconv"
	"time"
)

type SqlLogger struct {
	logger.Interface
}

func (l SqlLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
	sql, rows := fc()
	// 从上下文中获取分片信息
	if shard, ok := ctx.Value("shard_key").(string); ok {
		fmt.Printf("Shard: %s, SQL: %s, RowsAffected: %d\n", shard, sql, rows)
	} else {
		fmt.Printf("SQL: %s, RowsAffected: %d\n", sql, rows)
	}
}

// 只分表

func RunTbsharding() {
	db := initDb()

	//db.Use(sharding.Register(sharding.Config{
	//	ShardingKey:         "order_id",
	//	NumberOfShards:      3,
	//	PrimaryKeyGenerator: sharding.PKSnowflake,
	//}, "t_order"))

	// 自定义分表策略
	db.Use(sharding.Register(sharding.Config{
		ShardingKey: "order_id",
		ShardingAlgorithm: func(columnValue any) (suffix string, err error) {
			fmt.Println("----------", columnValue)
			return strconv.Itoa(columnValue.(int) % 3), nil
		},
	}, "t_order"))

	o := entity.Order{}
	db.Raw("select * from t_order where order_id = ? limit 1", 228620604381593600).Scan(&o)
	fmt.Printf("-----------%+v", o)

}

func initDb() *gorm.DB {
	dsn := "root:fan@1988fan@1988fan@1988fan@1988@tcp(127.0.0.1:3306)/sharding_jdbc0?charset=utf8mb4&parseTime=True&loc=Local"
	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
		//Logger: logger.Default.LogMode(logger.Info), //配置日志级别，打印出所有的sql
		Logger: SqlLogger{},
		NamingStrategy: schema.NamingStrategy{
			TablePrefix: "t_", // 表名前缀，`User` 的表名应该是 `t_users`
		},
	})
	if err != nil {
		fmt.Println("open mysql failed,", err)
		panic(err)
	}

	// 注册插件
	//db.Use()
	//db.Use(crypto.NewCryptoPlugin())

	sqlDB, _ := db.DB()
	sqlDB.SetMaxIdleConns(10)
	sqlDB.SetMaxOpenConns(500)
	sqlDB.SetConnMaxLifetime(time.Hour)

	return db
}
